import socket
import threading
import base64
import hashlib
import struct

# 1. 创建一个TCP服务器
# 2. 监听客户端连接
# 3. 接收客户端数据
# 4. 处理客户端数据
# 5. 发送响应数据
# 6. 关闭连接   

class TCPServer:
    def __init__(self, host='127.0.0.1', port=5000):
        self.host = host
        self.port = port
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.bind((self.host, self.port))
        self.server_socket.listen(5)
        print(f"Server listening on {self.host}:{self.port}")

    def handle_client(self, client_socket, addr):
        print(f"New connection from {addr}")
        is_websocket = False
        while True:
            try:
                data = client_socket.recv(1024)
                if not data:
                    break
                
                # Check for WebSocket handshake request
                if b"Upgrade: websocket" in data and not is_websocket:
                    print(f"WebSocket handshake request received from {addr}")
                    # Parse WebSocket key
                    key = None
                    for line in data.split(b'\r\n'):
                        if b'Sec-WebSocket-Key:' in line:
                            key = line.split(b': ')[1].strip()
                            print(f"WebSocket key found: {key}")
                            break
                    
                    if key:
                        # Calculate WebSocket accept key
                        GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
                        accept = base64.b64encode(
                            hashlib.sha1(key + GUID.encode('utf-8')).digest()
                        ).decode('utf-8')
                        print(f"Calculated Sec-WebSocket-Accept: {accept}")
                        
                        # Send WebSocket handshake response
                        response = (
                            b"HTTP/1.1 101 Switching Protocols\r\n"
                            b"Upgrade: websocket\r\n"
                            b"Connection: Upgrade\r\n"
                            b"Sec-WebSocket-Accept: " + accept.encode('utf-8') + b"\r\n\r\n"
                        )
                        client_socket.send(response)
                        print(f"WebSocket handshake completed with {addr}")
                        is_websocket = True
                        continue
                    else:
                        print("WebSocket key not found, closing connection.")
                        break
                
                if is_websocket:
                    # Handle WebSocket data frame
                    byte1, byte2 = data[:2]
                    opcode = byte1 & 0b00001111
                    if opcode == 0x8:  # Connection close frame
                        print(f"Received close frame from {addr}")
                        break
                    payload_length = byte2 & 0b01111111
                    if payload_length == 126:
                        payload_length = struct.unpack(">H", data[2:4])[0]
                        masks = data[4:8]
                        payload = data[8:8+payload_length]
                    elif payload_length == 127:
                        payload_length = struct.unpack(">Q", data[2:10])[0]
                        masks = data[10:14]
                        payload = data[14:14+payload_length]
                    else:
                        masks = data[2:6]
                        payload = data[6:6+payload_length]
                    
                    decoded_data = bytearray(b ^ masks[i % 4] for i, b in enumerate(payload))
                    try:
                        message = decoded_data.decode('utf-8')
                        print(f"Received WebSocket data from {addr}: {message}")
                        response = f"Server received: {message}"
                        # Send back a WebSocket text frame
                        response_frame = b'\x81' + struct.pack('B', len(response)) + response.encode('utf-8')
                        client_socket.send(response_frame)
                    except UnicodeDecodeError:
                        print(f"Failed to decode WebSocket data from {addr}: {decoded_data}")
                        break
                else:
                    # Handle regular data
                    try:
                        decoded_data = data.decode('utf-8')
                        print(f"Received from {addr}: {decoded_data}")
                        response = f"Server received: {decoded_data}"
                    except UnicodeDecodeError:
                        print(f"Received binary data from {addr}: {data!r}")
                        response = f"Server received binary data: {len(data)} bytes"
                    
                    client_socket.send(response.encode('utf-8'))
            except Exception as e:
                print(f"Error handling client {addr}: {e}")
                break
        client_socket.close()
        print(f"Connection from {addr} closed")

    def start(self):
        while True:
            client_socket, addr = self.server_socket.accept()
            # 为每个客户端创建一个新线程
            client_thread = threading.Thread(
                target=self.handle_client,
                args=(client_socket, addr)
            )
            client_thread.start()

if __name__ == "__main__":
    server = TCPServer()
    server.start()
